"""
使用sklearn的API实现线性回归
"""
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import sklearn.linear_model as lm # 线性模型模块
import sklearn.metrics as sm # 模型评估模块
import pickle # 模型保存

from pandas.core.interchange.from_dataframe import primitive_column_to_ndarray

# 数据准备
data = pd.read_csv('Salary_Data.csv')
# print(data)

# 绘制散点图，查看数据分布
# plt.scatter(data['YearsExperience'],data['Salary'])
# plt.show()

# 整理输入数据（二维），输出数据（一维）
train_x = data.iloc[:,:-1] #  不包含最后一列
# print("train_x")
# print(train_x)
train_y = data.iloc[:,-1] # 只包含最后一列
# print("train_y")
# print(train_y)

# 构建模型
model = lm.LinearRegression()
# 训练模型
model.fit(train_x,train_y)

# 测试模型
pred_train_y = model.predict(train_x)

# plt.plot(train_x,pred_train_y,c='orangered')
# # 绘制散点图，查看数据分布
# plt.scatter(data['YearsExperience'],data['Salary'])
# plt.show()

# 拿到一组测试集
# 假设我们的测试数据，没有参加过训练
test_x = train_x.iloc[::4]
test_y = train_y[::4]
# 预测值
pred_test_y = model.predict(test_x)
# 模型评估指标
# 平均绝对值误差 mae
print(sm.mean_absolute_error(test_y,pred_test_y))
# 平均平方误差：均方误差 mse
print(sm.mean_squared_error(test_y,pred_test_y))
# 中位数绝对偏差
print(sm.median_absolute_error(test_y,pred_test_y))
# R2_score
print(sm.r2_score(test_y,pred_test_y))

# 模型的保存
with open('model.pickle','wb') as f:
    pickle.dump(model,f)
print('模型保存成功')

# 模型的加载
with open('model.pickle','rb') as f:
    new_model = pickle.load(f)
print(new_model.predict([[1.1]]))




